#include "diskinfo.h"
#include <cerrno>
#include <stdio.h>
#include <stdlib.h>
#include <fstream>
#include <sstream>
#include <cstring>
#if defined(WIN32) || defined(WIN64) || defined(_WIN32) || defined(_WIN64)
#include <windows.h>
#include <fileapi.h>
#include <winioctl.h>
#else
#include <unistd.h>
#include <fcntl.h>
#include <dirent.h>
#include <sys/statvfs.h>
#include <sys/ioctl.h>
#include <scsi/scsi.h>
#include <scsi/sg.h>
#include <sys/sysmacros.h>
#endif



#if defined(WIN32) || defined(WIN64) || defined(_WIN32) || defined(_WIN64)
// Windows实现获取磁盘信息
static void getWindowsDiskInfo(std::vector<DiskInfo>& vtDiskInfos)
{
    std::string strPath = "C://";
    ULARGE_INTEGER totalBytes, freeBytes;
    if (GetDiskFreeSpaceExA(strPath.c_str(), nullptr, &totalBytes, &freeBytes)) {
        DiskInfo info;

        info.dTotalKB = totalBytes.QuadPart / 1024.0;
        info.dAvailKB = freeBytes.QuadPart / 1024.0;

        std::string driveType = "Unknown";
        UINT type = GetDriveTypeA(strPath.c_str());
        switch (type) {
        case DRIVE_FIXED:       driveType = "Fixed Disk"; break;
#if 0
        case DRIVE_SSD:         driveType = "SSD";         break;
#endif
        case DRIVE_REMOTE:      driveType = "Network Drive"; break;
        case DRIVE_CDROM:       driveType = "Optical Drive"; break;
        case DRIVE_REMOVABLE:   driveType = "Removable Drive"; break;
        }

        info.strType = driveType;

        // TODO 其他业务待实现
        (void)vtDiskInfos;
    }

    return;
}
#else
// Linux实现获取磁盘信息
static void getLinuxDiskInfo(std::vector<DiskInfo>& vtDiskInfos)
{
    // 读取/sys/block下的设备
    const char* const pDir = "/sys/block";
    DIR* dir = opendir(pDir);
    if (!dir) {
        printf(">>> Open directory (%s) failed.\r\n", pDir);
        return;
    }

    struct dirent* entry = nullptr;
    while ((entry = readdir(dir)) != nullptr) {
        std::string name = entry->d_name;

        // 跳过虚拟设备
        if (name == "." || name == ".."
                || name.find("loop") == 0 || name.find("ram") == 0
                || name.find("md0") == 0 || name.find("sr") == 0) {
            continue;
        }

        DiskInfo disk;
        disk.strName = name;

        // 尝试获取型号和序列号
        std::string modelPath = "/sys/block/" + name + "/device/model";
        std::ifstream modelFile(modelPath);
        if (modelFile) {
            getline(modelFile, disk.strModel);
        }

        std::string serialPath = "/sys/block/" + name + "/device/serial";
        std::ifstream serialFile(serialPath);
        if (serialFile) {
            std::getline(serialFile, disk.strSerial);
        }

        // 通过/sys/block判断设备类型
        std::ifstream rotational("/sys/block/" + name + "/queue/rotational");
        disk.strType = rotational ? "HDD" : "SSD";
#if 1
        // 简化的健康状态检测
        std::string statPath = "/sys/block/" + name + "/stat";
        std::ifstream statFile(statPath);
        if (statFile) {
            // 读取统计信息
            long readIOs, readMerges, readSectors, readTicks;
            long writeIOs, writeMerges, writeSectors, writeTicks;
            long inFlight, ioTicks, timeInQueue;

            statFile >> readIOs >> readMerges >> readSectors >> readTicks
                    >> writeIOs >> writeMerges >> writeSectors >> writeTicks
                    >> inFlight >> ioTicks >> timeInQueue;

            // 根据IO错误率评估健康状态
            if (readIOs > 1000000 && (readIOs / (readIOs + writeIOs) < 0.1)) {
                disk.strHealthStatus = "Warning";
            } else {
                disk.strHealthStatus = "Healthy";
            }
        } else {
            disk.strHealthStatus = "Unknown";
        }
#else
        // 读取错误计数
        std::ifstream errors("/sys/block/" + name + "/stat");
        std::string stats;
        std::getline(errors, stats);
        size_t readErrors = stats.find(' ') != std::string::npos ? std::stoul(stats) : 0;
        disk.healthStatus = disk.type + " | Read Errors: " + std::to_string(readErrors);
#endif
        // 检查是否存在分区，并进行分区处理
        std::string basePath = "/sys/block/" + name;
        DIR* subdir = opendir(basePath.data());
        if (subdir) {
            struct dirent* subEntry = nullptr;
            while ((subEntry = readdir(subdir)) != nullptr) {
                std::string subName = subEntry->d_name;
                if (subName.find(name) == 0 && subName.size() > name.size()) {
                    // 匹配分区的磁盘，如sda1，sda2等
                    disk.strName = subName; // 实际的名称

                    // 获取挂载的path
                    std::ifstream mounts("/proc/mounts");
                    std::string line = "";

                    while (std::getline(mounts, line)) {
                        std::istringstream iss(line);
                        std::string mountDevice, mountPoint, fsType;
                        iss >> mountDevice >> mountPoint >> fsType;

                        if (mountDevice == "/dev/" + disk.strName) {
                            // 被真实的挂载
                            disk.strPath = mountPoint;

                            struct statvfs buf;
                            if (statvfs(disk.strPath.c_str(), &buf) == 0) {
                                disk.dTotalKB = (buf.f_blocks * buf.f_frsize) / 1024.0;
                                disk.dAvailKB = (buf.f_bfree * buf.f_frsize) / 1024.0;
                            }

                            vtDiskInfos.emplace_back(disk);
                            break;
                        }
                    }
                }
            }

            closedir(subdir);
        }
    }

    closedir(dir);

    return;
}
#endif

/********************************************************
 * 功能：获取系统磁盘信息
 * 输出参数：
 *      vtDiskInfos：磁盘信息集合
 * ******************************************************/
void getDiskInfo(std::vector<DiskInfo>& vtDiskInfos)
{
#if defined(WIN32) || defined(WIN64) || defined(_WIN32) || defined(_WIN64)
    getWindowsDiskInfo(vtDiskInfos);
#else
    getLinuxDiskInfo(vtDiskInfos);
#endif
    return;
}

